package com.flame.ai.pdf.transformer;

import com.flame.ai.pdf.bean.ImageItem;
import com.flame.ai.pdf.bean.TableItem;
import com.flame.ai.pdf.bean.TextItem;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * @author Jflame
 * @version 1.0.0
 * @since 2025/8/27
 */
public class DocumentTransformer {

    public static String DEFAULT_IMAGE_DESCRIPTION_TEMPLATE = """
			Here is the context of the image:
			{context_str}
			
			Here is the caption of the image:
			{caption_str}
			
			Here is the footnote of the image:
			{footnote_str}

			Description:""";

    public static String DEFAULT_TABLE_DESCRIPTION_TEMPLATE = """
			Here is the context of the table:
			{context_str}
			
			Here is the caption of the table:
			{caption_str}
			
			Here is the footnote of the table:
			{footnote_str}

			Description:""";

	private static final String FOOTNOTE_STR_PLACEHOLDER = "footnote_str";

	private static final String CAPTION_STR_PLACEHOLDER = "caption_str";

	private static final String CONTEXT_STR_PLACEHOLDER = "context_str";


	private final ChatModel chatModel;

	public DocumentTransformer(String multiModel) {
		OllamaOptions ollamaOptions = OllamaOptions.builder()
				.model(multiModel)
				.temperature(0.0)
				.build();
		this.chatModel = OllamaChatModel.builder()
				.ollamaApi(OllamaApi.builder().build())
				.defaultOptions(ollamaOptions)
				.build();
	}

	public List<Document> transformImages(List<ImageItem> imageItems) {
		ArrayList<Document> documents = new ArrayList<>();

		for (ImageItem imageItem : imageItems) {

			String caption = imageItem.getCaption() != null && !imageItem.getCaption().isEmpty() ? String.join(" ", imageItem.getCaption()) : "null";
			String footnote = imageItem.getFootnote() != null && !imageItem.getFootnote().isEmpty() ? String.join(" ", imageItem.getFootnote()) : "null";

			Prompt prompt = new PromptTemplate(DEFAULT_IMAGE_DESCRIPTION_TEMPLATE)
					.create(Map.of(CONTEXT_STR_PLACEHOLDER, imageItem.getContext(),
							FOOTNOTE_STR_PLACEHOLDER, footnote,
							CAPTION_STR_PLACEHOLDER,  caption));

			documents.add(new Document(
					Objects.requireNonNull(chatModel.call(prompt).getResult().getOutput().getText()),
					imageItem.getMetadata()));
		}
		return documents;
	}

	public List<Document> transformTables(List<TableItem> tableItems) {
		ArrayList<Document> documents = new ArrayList<>();

		for (TableItem tableItem : tableItems) {

			String caption = tableItem.getCaption() != null && !tableItem.getCaption().isEmpty() ? String.join(" ", tableItem.getCaption()) : "null";
			String footnote = tableItem.getFootnote() != null && !tableItem.getFootnote().isEmpty() ? String.join(" ", tableItem.getFootnote()) : "null";

			Prompt prompt = new PromptTemplate(DEFAULT_TABLE_DESCRIPTION_TEMPLATE)
					.create(Map.of(CONTEXT_STR_PLACEHOLDER, tableItem.getContext(),
							FOOTNOTE_STR_PLACEHOLDER, footnote,
							CAPTION_STR_PLACEHOLDER,  caption));

			documents.add(new Document(
					Objects.requireNonNull(chatModel.call(prompt).getResult().getOutput().getText()),
					tableItem.getMetadata()));
		}
		return documents;
	}

	public List<Document> transformTexts(List<TextItem> textItems) {
		ArrayList<Document> documents = new ArrayList<>();
		for (TextItem textItem : textItems) {
			documents.add(new Document(textItem.getText(), textItem.getMetadata()));
		}
		return documents;
	}
}
